#!/usr/bin/python2.4


import _mysql_exceptions
import itertools
import operator
import MySQLdb
import pprint
import re
import time

import constants
import flags
from logger import logger, rpc_logger
import rpc_stubs
import sql_db_object
import sql_job_host_run
import utils


flags.add_option('--log-sql-select', '--lss', action='store_true')


class SqlQuery(object):

  def __init__(self):
    self.conn_dargs = {'host': rpc_stubs.AUTOTEST_WEB}
    for p in ('host', 'port', 'user', 'passwd'):
      v = getattr(constants, ('sql_' + p).upper(), None)
      if v:
        self.conn_dargs[p] = v
    self.conn = None
    self.cursor = None
    self.new_schema = None
    self.table_rename_re = None

  def __del__(self):
    if self.cursor:
      self.cursor.close()
    if self.conn:
      self.conn.close()

  def setup(self):
    self.conn = MySQLdb.connect(**self.conn_dargs)
    self.conn.autocommit(True)
    self.conn.converter[MySQLdb.constants.FIELD_TYPE.LONG] = int
    self.conn.converter[MySQLdb.constants.FIELD_TYPE.LONGLONG] = int
    self.conn.converter.pop(MySQLdb.constants.FIELD_TYPE.TIMESTAMP, None)
    self.conn.converter.pop(MySQLdb.constants.FIELD_TYPE.DATE, None)
    self.conn.converter.pop(MySQLdb.constants.FIELD_TYPE.TIME, None)
    self.conn.converter.pop(MySQLdb.constants.FIELD_TYPE.DATETIME, None)
    self.cursor = self.conn.cursor()
    self.conn.select_db('autotest_web')
    self.new_schema = self.cursor.execute('SHOW TABLES LIKE "tko_%"') > 0
    if self.new_schema:
      self.table_rename_re = re.compile(
          r'\b((?:FROM|JOIN)\s+)(\w+)(\s+AS\s+\w+)\b')

  def handle_new_schema(self, db, statement):
    if self.new_schema and db:
      table_name_prefix = ''
      if db == 'autotest_web':
        table_name_prefix = 'afe_'
      elif db == 'tko':
        db = 'autotest_web'
        table_name_prefix = 'tko_'
      if table_name_prefix:
        replacement = r'\1%s\2\3' % table_name_prefix
        statement = self.table_rename_re.sub(replacement, statement)
    return db, statement

  def execute(self, statement, db=None, log_results=None, caller_level=4):
    if not self.cursor:
      self.setup()

    db, statement = self.handle_new_schema(db, statement)

    if flags.get_option('log_sql_select'):
      rpc_logger.debug('SQL use %s; %s' % (db, statement),
                       caller_level=caller_level)
    else:
      rpc_logger.debug('SQL use %s; SELECT ...' % db,
                       caller_level=caller_level)
    start_time = time.time()
    try:
      if db:
        self.conn.select_db(db)
      results = self.cursor.execute(statement)
    except _mysql_exceptions.OperationalError, e:
      if getattr(e, 'args', None) != (2006, 'MySQL server has gone away'):
        raise
      rpc_logger.info('MySQL server has gone away, reconnecting ...')
      self.setup()
      if db:
        self.conn.select_db(db)
      results = self.cursor.execute(statement)
    sql_finish_time = time.time()
    rpc_logger.debug('columns: %s' %
                     ', '.join(col[0] for col in self.cursor.description),
                     caller_level=caller_level)
    if log_results:
      rpc_logger.debug(pprint.pformat([col for col in self.cursor]),
                       caller_level=caller_level)
    log_finish_time = time.time()
    rpc_logger.debug('%d rows in set (%.2g seconds),'
                     ' logging done in %.2g seconds' %
                     (results,
                      sql_finish_time - start_time,
                      log_finish_time - sql_finish_time),
                     caller_level=caller_level)
    return results

  def log_query(self, name, **filter_data):
    rpc_logger.debug('SQL sql.%s(%s)' % (
        name, ', '.join('%s=%r' % kv
                        for kv in filter_data.iteritems())))

  def get_hosts_use_algo(self, algo, **filter_data):
    allowed_algos = ('count', 'join')
    if algo not in allowed_algos:
      logger.error('Arg algo must be one of %r' % (allowed_algos,))
      return None

    log_results = filter_data.pop('log_results', None)
    sort_by = filter_data.pop('sort_by', ())
    multiple_labels = filter_data.pop('multiple_labels', ())

    fields = ('hosts.id', 'hosts.hostname', 'hosts.locked',
              'hosts.status', 'hosts.protection',
              'GROUP_CONCAT(DISTINCT IF(labels.platform, labels.name, NULL))',
              'GROUP_CONCAT(DISTINCT labels.name)',
              'ISNULL(host_queue_entries.host_id)'
              ' AND ISNULL(special_tasks.host_id)',
              'host_attributes.value')

    select = 'SELECT '
    select += ','.join('%s AS %s' % (field, key)
                       for field, key in itertools.izip(
                           fields, sql_db_object.Host.KEYS))

    select += ' FROM hosts AS hosts'

    if algo == 'join':
      if multiple_labels:
        for i, label in enumerate(multiple_labels):
          if multiple_labels.index(label) == i:
            select += (
                ' INNER JOIN hosts_labels AS hosts_labels_%(i)s'
                ' ON hosts.id = hosts_labels_%(i)s.host_id'
                ' INNER JOIN labels AS labels_%(i)s'
                ' ON hosts_labels_%(i)s.label_id = labels_%(i)s.id'
                ' AND labels_%(i)s.name = %(label)r'
                ' AND NOT labels_%(i)s.invalid' %
                dict(i=i, label=label))

    select += (' LEFT JOIN hosts_labels AS hosts_labels'
               ' ON hosts.id = hosts_labels.host_id'
               ' LEFT JOIN labels AS labels'
               ' ON hosts_labels.label_id = labels.id')
    select += (' LEFT JOIN host_queue_entries AS host_queue_entries'
               ' ON hosts.id = host_queue_entries.host_id'
               ' AND NOT host_queue_entries.complete')
    select += (' LEFT JOIN special_tasks AS special_tasks'
               ' ON hosts.id = special_tasks.host_id'
               ' AND NOT special_tasks.is_complete')
    select += (' LEFT JOIN host_attributes AS host_attributes'
               ' ON hosts.id = host_attributes.host_id'
               ' AND host_attributes.attribute = %r' %
               constants.HA_NOTES)

    select += (' WHERE NOT hosts.invalid AND'
               ' (NOT labels.invalid OR hosts_labels.host_id IS NULL)')

    having_filters = []
    if filter_data:
      name_to_expr = dict(itertools.izip(sql_db_object.Host.KEYS, fields))
      for name, op, value in utils.parse_filters(**filter_data):
        expr = name_to_expr.get(name, name)
        if expr.startswith('GROUP_CONCAT('):
          having_filters.append(utils.convert_to_sql(name, op, value))
        else:
          select += ' AND %s' % utils.convert_to_sql(expr, op, value)

    select += ' GROUP BY hosts.id'

    having_added = False

    if algo == 'count':
      if multiple_labels:
        select += (' HAVING COUNT(DISTINCT'
                   ' IF(labels.name IN (%s), labels.name, NULL)) = %d' %
                   (','.join(repr(label) for label in set(multiple_labels)),
                    len(set(multiple_labels))))
        having_added = True

    if having_filters:
      if having_added:
        select += ' AND '
      else:
        select += ' HAVING '
        having_added = True
      select += ' AND '.join(having_filters)

    self.execute(select, 'autotest_web', log_results=log_results)
    hosts = [sql_db_object.Host(raw_object)
             for raw_object in self.cursor]

    if (len(sort_by) > 1 or
        (len(sort_by) > 0 and sort_by[0] != 'id')):
      utils.sort_dicts(hosts, sort_by)

    return hosts

  def get_hosts(self, **filter_data):
    self.log_query('get_hosts', **filter_data)
    return self.get_hosts_use_algo(constants.GET_SQL_HOSTS_USE_ALGO,
                                   **filter_data)

  def get_test_views_use_algo(self, algo, **filter_data):
    allowed_algos = ('count', 'join', 'exists')
    if algo not in allowed_algos:
      logger.error('Arg algo must be one of %r' % (allowed_algos,))
      return None

    log_results = filter_data.pop('log_results', None)
    sort_by = filter_data.pop('sort_by', ())
    include_labels_set = set(filter_data.pop('include_labels', ()))
    exclude_labels_set = set(filter_data.pop('exclude_labels', ()))

    fields = tuple('test_view_2.%s' % key
                   for key in sql_db_object.TestView.KEYS[:-2])
    fields += ('test_attributes.value',
               'GROUP_CONCAT(DISTINCT test_labels.name)')

    select = 'SELECT '
    select += ','.join('%s AS %s' % (field, key)
                       for field, key in itertools.izip(
                           fields, sql_db_object.TestView.KEYS))

    select += ' FROM test_view_2 AS test_view_2'

    if algo == 'join':
      if include_labels_set:
        select += (
            ' INNER JOIN test_labels_tests AS test_labels_tests_include'
            ' ON test_view_2.test_idx = test_labels_tests_include.test_id'
            ' AND test_labels_tests_include.testlabel_id IN'
            ' (SELECT id FROM test_labels AS test_labels WHERE name IN (%s))' %
            ','.join(repr(label) for label in include_labels_set))

      if exclude_labels_set:
        select += (
            ' LEFT JOIN test_labels_tests AS test_labels_tests_exclude'
            ' ON test_view_2.test_idx = test_labels_tests_exclude.test_id'
            ' AND test_labels_tests_exclude.testlabel_id IN'
            ' (SELECT id FROM test_labels AS test_labels WHERE name in (%s))' %
            ','.join(repr(label) for label in exclude_labels_set))
        join_exclude_where = 'test_labels_tests_exclude.test_id IS NULL'
      else:
        join_exclude_where = ''

    select += (' LEFT JOIN test_labels_tests AS test_labels_test'
               ' ON test_view_2.test_idx = test_labels_tests.test_id'
               ' LEFT JOIN test_labels AS test_labels'
               ' ON test_labels_tests.testlabel_id = test_labels.id')

    select += (' LEFT JOIN test_attributes AS test_attributes'
               ' ON test_view_2.test_idx = test_attributes.test_idx'
               ' AND test_attributes.attribute = %r' %
               constants.TA_TRIAGE_INFO)

    where_added = False

    if algo == 'join' and join_exclude_where:
      select += ' WHERE %s' % join_exclude_where
      where_added = True

    if algo == 'exists':
      if include_labels_set:
        select += (' WHERE EXISTS (SELECT 1'
                   ' FROM test_labels_tests AS test_labels_test'
                   ' INNER JOIN test_labels AS test_labels'
                   ' ON test_labels_tests.testlabel_id = test_labels.id'
                   ' AND test_labels.name IN (%s)'
                   ' WHERE test_labels_tests.test_id = test_view_2.test_idx)' %
                   ','.join(repr(label) for label in include_labels_set))
        where_added = True

      if exclude_labels_set:
        if where_added:
          select += ' AND'
        else:
          select += ' WHERE'
          where_added = True
        select += (' NOT EXISTS (SELECT 1'
                   ' FROM test_labels_tests AS test_labels_tests'
                   ' INNER JOIN test_labels AS test_labels'
                   ' ON test_labels_tests.testlabel_id = test_labels.id'
                   ' AND test_labels.name IN (%s)'
                   ' WHERE test_labels_tests.test_id = test_view_2.test_idx)' %
                   ','.join(repr(label) for label in exclude_labels_set))

    having_filters = []
    if filter_data:
      name_to_expr = dict(itertools.izip(sql_db_object.TestView.KEYS, fields))
      for name, op, value in utils.parse_filters(**filter_data):
        expr = name_to_expr.get(name, name)
        if expr.startswith('GROUP_CONCAT('):
          having_filters.append(utils.convert_to_sql(name, op, value))
        else:
          if where_added:
            select += ' AND '
          else:
            select += ' WHERE '
            where_added = True
          select += utils.convert_to_sql(expr, op, value)

    select += ' GROUP BY test_view_2.job_idx,test_view_2.test_idx'

    having_added = False

    if algo == 'count':
      if include_labels_set:
        select += (' HAVING COUNT(IF(test_labels.name IN (%s),'
                   ' test_labels.name, NULL))' %
                   ','.join(repr(label) for label in include_labels_set))
        having_added = True

      if exclude_labels_set:
        if having_added:
          select += ' AND '
        else:
          select += ' HAVING '
          having_added = True
        select += ('NOT COUNT(IF(test_labels.name IN (%s),'
                   ' test_labels.name, NULL))' %
                   ','.join(repr(label) for label in exclude_labels_set))

    if having_filters:
      if having_added:
        select += ' AND '
      else:
        select += ' HAVING '
        having_added = True
      select += ' AND '.join(having_filters)

    self.execute(select, 'tko', log_results=log_results)
    test_views = [sql_db_object.TestView(raw_object)
                  for raw_object in self.cursor]

    if (len(sort_by) > 2 or
        (len(sort_by) > 1 and sort_by[1] != 'test_idx') or
        (len(sort_by) > 0 and sort_by[0] != 'job_idx')):
      utils.sort_dicts(test_views, sort_by)

    return test_views

  def get_test_views(self, **filter_data):
    self.log_query('get_test_views', **filter_data)
    return self.get_test_views_use_algo(constants.GET_SQL_TEST_VIEWS_USE_ALGO,
                                        **filter_data)

  def get_jobs(self, **filter_data):
    log_results = filter_data.pop('log_results', None)
    sort_by = filter_data.pop('sort_by', ())
    include_labels_set = set(filter_data.pop('include_labels', ()))
    exclude_labels_set = set(filter_data.pop('exclude_labels', ()))

    fields = ('jobs.job_idx', 'jobs.tag', 'jobs.label', 'jobs.username',
              'machines.hostname',
              'jobs.queued_time', 'jobs.started_time', 'jobs.finished_time',
              "GROUP_CONCAT(DISTINCT test_attributes.value SEPARATOR '; ')",
              'GROUP_CONCAT(DISTINCT test_labels.name)')

    select = 'SELECT '
    select += ','.join('%s AS %s' % (field, key)
                       for field, key in itertools.izip(
                           fields, sql_db_object.Job.KEYS))

    select += ' FROM jobs AS jobs'

    select += (' LEFT JOIN machines AS machines'
               ' ON jobs.machine_idx = machines.machine_idx'
               ' LEFT JOIN tests AS tests'
               ' ON jobs.job_idx = tests.job_idx')

    select += (' LEFT JOIN test_attributes AS test_attributes'
               ' ON tests.test_idx = test_attributes.test_idx'
               ' AND test_attributes.attribute = %r' %
               constants.TA_TRIAGE_INFO)

    select += (' LEFT JOIN test_labels_tests AS test_labels_tests'
               ' ON tests.test_idx = test_labels_tests.test_id'
               ' LEFT JOIN test_labels AS test_labels'
               ' ON test_labels_tests.testlabel_id = test_labels.id')

    where_added = False

    if include_labels_set:
      select += (' WHERE EXISTS (SELECT 1'
                 ' FROM tests AS tests'
                 ' INNER JOIN test_labels_tests AS test_labels_tests'
                 ' ON tests.test_idx = test_labels_tests.test_id'
                 ' INNER JOIN test_labels AS test_labels'
                 ' ON test_labels_tests.testlabel_id = test_labels.id'
                 ' AND test_labels.name IN (%s)'
                 ' WHERE tests.job_idx = jobs.job_idx)' %
                   ','.join(repr(label) for label in include_labels_set))
      where_added = True

    if exclude_labels_set:
      if where_added:
        select += ' AND'
      else:
        select += ' WHERE'
        where_added = True
      select += (' NOT EXISTS (SELECT 1'
                 ' FROM tests AS tests'
                 ' INNER JOIN test_labels_tests AS test_labels_tests'
                 ' ON tests.test_idx = test_labels_tests.test_id'
                 ' INNER JOIN test_labels AS test_labels'
                 ' ON test_labels_tests.testlabel_id = test_labels.id'
                 ' AND test_labels.name IN (%s)'
                 ' WHERE tests.job_idx = jobs.job_idx)' %
                   ','.join(repr(label) for label in exclude_labels_set))

    having_filters = []
    if filter_data:
      name_to_expr = dict(itertools.izip(sql_db_object.Job.KEYS, fields))
      for name, op, value in utils.parse_filters(**filter_data):
        expr = name_to_expr.get(name, name)
        if expr.startswith('GROUP_CONCAT('):
          having_filters.append(utils.convert_to_sql(name, op, value))
        else:
          if where_added:
            select += ' AND '
          else:
            select += ' WHERE '
            where_added = True
          select += utils.convert_to_sql(expr, op, value)

    select += ' GROUP BY jobs.job_idx'

    if having_filters:
      select += ' HAVING '
      select += ' AND '.join(having_filters)

    self.execute(select, 'tko', log_results=log_results)
    jobs = [sql_db_object.Job(raw_object)
            for raw_object in self.cursor]

    if (len(sort_by) > 1 or
        (len(sort_by) > 0 and sort_by[0] != 'job_idx')):
      utils.sort_dicts(jobs, sort_by)

    return jobs

  def get_tests(self, **filter_data):
    log_results = filter_data.pop('log_results', None)
    sort_by = filter_data.pop('sort_by', ())

    fields = ('tests.test_idx', 'tests.job_idx', 'tests.test', 'status.word',
              'tests.reason', 'tests.started_time', 'tests.finished_time')

    select = 'SELECT '
    select += ','.join('%s AS %s' % (field, key)
                       for field, key in itertools.izip(
                           fields, sql_db_object.Test.KEYS))

    select += ' FROM tests AS tests'
    select += ' LEFT JOIN status AS status on tests.status = status.status_idx'

    if filter_data:
      name_to_expr = dict(itertools.izip(sql_db_object.Test.KEYS, fields))
      select += ' WHERE '
      select += ' AND '.join(
          utils.convert_to_sql(name_to_expr.get(name, name), op, value)
          for name, op, value in utils.parse_filters(**filter_data))

    if sort_by:
      select += ' ORDER BY '
      select += ','.join(sort_by)

    self.execute(select, 'tko', log_results=log_results)
    tests = [sql_db_object.Test(raw_object)
             for raw_object in self.cursor]

    return tests

  def get_job_host_runs(self, **filter_data):
    self.log_query('get_job_host_runs', **filter_data)

    sort_by = filter_data.pop('sort_by', ())

    job_sort_by = []
    test_sort_by = []
    for key in sort_by:
      if key in sql_db_object.Job.KEYS:
        job_sort_by.append(key)
      else:
        test_sort_by.append(key)

    job_filter_data = {}
    test_filter_data = {}
    for key, value in filter_data.iteritems():
      if key == 'log_results':
        job_filter_data[key] = test_filter_data[key] = value
      elif (key.rsplit('__', 1)[0] in sql_db_object.Job.KEYS or
          key in ('include_labels', 'exclude_labels')):
        job_filter_data[key] = value
      else:
        test_filter_data[key] = value

    self.execute('SELECT max(job_idx) FROM jobs AS jobs', 'tko',
                 log_results=filter_data.get('log_results'))
    max_job_idx = self.cursor.fetchone()[0]
    job_filter_data['job_idx__lte'] = min(
        max_job_idx, job_filter_data.get('job_idx__lte', max_job_idx))

    jobs = self.get_jobs(**job_filter_data)
    if jobs:
      tests = self.get_tests(sort_by=['job_idx'] + test_sort_by,
                             job_idx__in=(job.job_idx for job in jobs),
                             **test_filter_data)
    else:
      tests = []

    job_host_runs = []
    job_i = 0
    for job_idx, job_tests in itertools.groupby(
        tests, operator.attrgetter('job_idx')):
      while job_i < len(jobs) and jobs[job_i].job_idx < job_idx:
        if jobs[job_i].job_finished_time:
          logger.error('Job %r does not have any tests' % jobs[job_i].job_tag)
        job_host_runs.append(
            sql_job_host_run.JobHostRunWithJobTest(jobs[job_i], None))
        job_i += 1
      if job_i >= len(jobs) or jobs[job_i].job_idx > job_idx:
        logger.fatal('Unable to find job for tests with job_idx %s' % job_idx)
      job_host_runs.append(
          sql_job_host_run.JobHostRunWithJobTest(jobs[job_i], list(job_tests)))
      job_i += 1

    while job_i < len(jobs):
      if jobs[job_i].job_finished_time:
        logger.error('Job %r does not have any tests' % jobs[job_i].job_tag)
      job_host_runs.append(
          sql_job_host_run.JobHostRunWithJobTest(jobs[job_i], None))
      job_i += 1

    if (len(job_sort_by) > 1 or
        (len(job_sort_by) > 0 and job_sort_by[0] != 'job_idx')):
      utils.sort_objects(job_host_runs, sort_by)

    return job_host_runs


sql = SqlQuery()
